package com.dts.gateway.util;

import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * （自定义类描述）
 *
 * @author zhoujh
 * @date 2022/06/02 14:33
 */
public class RJLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    private final Logger log = LoggerFactory.getLogger(RJLoadBalancer.class);
    //网关所在区域
    private final String gatewayZone;
    //转发服务的名称
    private final String serviceId;
    //轮询随机数
    private final AtomicInteger position;
    //区域轮询随机数
    private final AtomicInteger zonePosition;
    private final String CLIENT_ZONE_KEY = "service-zone";
    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public RJLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId, String gatewayZone) {
        this(serviceInstanceListSupplierProvider, serviceId, gatewayZone, new Random().nextInt(1000), new Random().nextInt(1000));
    }

    public RJLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId, String gatewayZone, int seedPosition, int zoneSeedPosition) {
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.gatewayZone = gatewayZone;
        this.serviceId = serviceId;
        this.position = new AtomicInteger(seedPosition);
        this.zonePosition = new AtomicInteger(zoneSeedPosition);
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next()
                .map(serviceInstances -> processInstanceResponse(supplier, serviceInstances));
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier, List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = this.getInstanceResponse(serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("Gateway No servers available:" + serviceId);
            }
            return new EmptyResponse();
        } else {
//            List<ServiceInstance> zoneHitService = instances.parallelStream().filter((service) -> {
//                String zone = service.getMetadata().get(CLIENT_ZONE_KEY);
//                if (gatewayZone.equals(zone)) {
//                    return true;
//                } else {
//                    return false;
//                }
//            }).collect(Collectors.toList());

            List<ServiceInstance> zoneHitService = instances.parallelStream().collect(Collectors.toList());
            if (log.isWarnEnabled()) {
                log.warn("Gateway load balance zone hit:" + JSON.toJSONString(zoneHitService));
            }
            if (zoneHitService.isEmpty()) {
                return roundRobin(position, instances);
            } else {
                return roundRobin(zonePosition, zoneHitService);
            }
        }
    }

    private Response<ServiceInstance> roundRobin(AtomicInteger position, List<ServiceInstance> instances) {
        int pos = Math.abs(position.incrementAndGet());
        return new DefaultResponse(instances.get(pos % instances.size()));
    }
}
